package com.tyf.common.zk.lock;

import com.tyf.common.exception.KeeperException;
import com.tyf.common.zk.ZKTemplete;
import com.tyf.common.zk.listener.KeeperNodeListener;
import org.apache.zookeeper.KeeperException.NoNodeException;
import org.apache.zookeeper.KeeperException.NodeExistsException;

import java.util.Set;
import java.util.concurrent.CopyOnWriteArraySet;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Lock;

/**
 * @author huangdou
 * @version 0.0.1
 * <p>
 * 同步锁，用于处理分布式多实例多线程之间的一致协作。
 * @at 2016年12月5日上午10:28:39
 */
public class KeeperCountDownLatch {
    public static final String LATCH_ROOT = "/latch_root";
    private String name;

    private ZKTemplete templete;

    private String latchPath;

    private Lock latchLock;

    private Set<Semaphore> waitSemaphore = new CopyOnWriteArraySet<Semaphore>();

    public KeeperCountDownLatch(String name, ZKTemplete templete, int count) {
        this.name = name;
        this.templete = templete;
        init(count);
    }

    public static KeeperCountDownLatch getOrCreate(String name,
                                                   int createCount, ZKTemplete templete) {
        if (name == null || "".equals(name.trim()) || name.contains("/")) {
            throw new IllegalArgumentException(String.format(
                    "name can not be %s", name));
        }
        if (createCount < 0) {
            throw new IllegalArgumentException(
                    String.format("count can not be less than 0"));
        }

        return new KeeperCountDownLatch(name, templete, createCount);
    }

    private void init(int count) {
        if (name == null || "".equals(name.trim()) || name.contains("/")) {
            throw new IllegalArgumentException(String.format(
                    "path can not be %s", name));
        }
        latchLock = new KeeperMutexLock((LATCH_ROOT + "_" + name).replace("/",
                ""), templete);
        latchPath = LATCH_ROOT + "/" + name;
        latchLock.lock();
        try {
            if (!templete.exist(latchPath)) {
                try {
                    templete.createWtihParent(latchPath);
                    templete.update(latchPath, (count + "").getBytes());
                } catch (KeeperException e) {
                    if (!(e.getProto() instanceof NodeExistsException)) {
                        throw e;
                    }
                }
            }
        } finally {
            latchLock.unlock();
        }
        templete.listenNode(latchPath, new LatchListener());
    }

    public boolean await(long timeout, TimeUnit unit) throws InterruptedException {
        Semaphore semaphore = new Semaphore(0);
        latchLock.lock();
        try {
            int count = Integer.parseInt(new String(templete.read(latchPath)));
            if (count == 0) {
                return true;
            } else {
                waitSemaphore.add(semaphore);
            }
        } catch (KeeperException e) {
            if (!(e.getProto() instanceof NoNodeException)) {
                throw new RuntimeException("Latch Has been destroyed!");
            }
        } finally {
            latchLock.unlock();
        }

        boolean isLocked = false;
        if (timeout > 0) {
            isLocked = semaphore.tryAcquire(timeout, unit);
        } else {
            isLocked = semaphore.tryAcquire();
        }
        return isLocked;
    }

    public void await() throws InterruptedException {
        Semaphore semaphore = new Semaphore(0);
        latchLock.lock();
        try {
            int count = Integer.parseInt(new String(templete.read(latchPath)));
            if (count == 0) {
                return;
            } else {
                waitSemaphore.add(semaphore);
            }
        } catch (KeeperException e) {
            if (!(e.getProto() instanceof NoNodeException)) {
                throw new RuntimeException("Latch Has been destroyed!");
            }
        } finally {
            latchLock.unlock();
        }
        semaphore.acquire();
    }

    public void destroy() {
        latchLock.lock();

        try {
            templete.delete(latchPath);
        } catch (KeeperException e) {
            if (!(e.getProto() instanceof NoNodeException)) {
                throw new RuntimeException("Latch Has been destroyed!");
            }
        } finally {
            latchLock.unlock();
        }

    }

    public void countDown() {
        latchLock.lock();

        try {
            int count = Integer.parseInt(new String(templete.read(latchPath)));
            if (count > 0) {
                templete.update(latchPath, ((--count) + "").getBytes());
            }
        } catch (KeeperException e) {
            if (!(e.getProto() instanceof NoNodeException)) {
                throw new RuntimeException("Latch Has been destroyed!");
            }
        } finally {
            latchLock.unlock();
        }
    }

    public long getCount() {
        try {
            return Integer.parseInt(new String(templete.read(latchPath)));
        } catch (KeeperException e) {
            if (!(e.getProto() instanceof NoNodeException)) {
                throw new RuntimeException("Latch Has been destroyed!");
            }
            throw e;
        }
    }

    class LatchListener implements KeeperNodeListener {
        @Override
        public void onData(String path, byte[] bytes) {
            if (bytes != null && Integer.parseInt(new String(bytes)) == 0) {
                for (Semaphore semaphore : waitSemaphore) {
                    semaphore.release();
                }
            }
        }

        @Override
        public void onDelete(String path) {
            for (Semaphore semaphore : waitSemaphore) {
                semaphore.release();
            }
        }

    }

}
